#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
传输优化工具类
支持多种传输方式和优化策略，提升内网传输速度
"""

import os
import gzip
import threading
import time
import logging
import subprocess
import tempfile
from concurrent.futures import ThreadPoolExecutor, as_completed
from pathlib import Path
import requests
from config import config

logger = logging.getLogger(__name__)

class TransferOptimizer:
    """传输优化器"""
    
    def __init__(self):
        self.config = config
        self.temp_dir = tempfile.gettempdir()
    
    def get_chunk_size(self):
        """动态获取切片大小"""
        return self.config.get_transfer_chunk_size()
    
    def get_use_compression(self):
        """动态获取是否使用压缩"""
        return self.config.get_transfer_use_compression()
    
    def get_compression_level(self):
        """动态获取压缩级别"""
        return self.config.get_transfer_compression_level()
    
    def get_use_parallel(self):
        """动态获取是否使用并行传输"""
        return self.config.get_transfer_use_parallel()
    
    def get_parallel_chunks(self):
        """动态获取并行传输块数"""
        return self.config.get_transfer_parallel_chunks()
    
    def get_use_http(self):
        """动态获取是否使用HTTP传输"""
        return self.config.get_transfer_use_http()
    
    def get_use_rsync(self):
        """动态获取是否使用rsync"""
        return self.config.get_transfer_use_rsync()
    
    def get_http_port(self):
        """动态获取HTTP传输端口"""
        return self.config.get_transfer_http_port()
    
    def optimize_image_save(self, image_name, output_path):
        """优化镜像保存过程"""
        try:
            logger.info(f"开始优化保存镜像: {image_name}")
            
            # 根据配置选择保存策略
            if self.get_use_compression():
                return self._save_with_compression(image_name, output_path)
            else:
                return self._save_without_compression(image_name, output_path)
                
        except Exception as e:
            logger.error(f"优化保存镜像失败: {str(e)}")
            return None, str(e)
    
    def _save_with_compression(self, image_name, output_path):
        """使用压缩保存镜像"""
        try:
            # 创建临时tar文件
            temp_tar = os.path.join(self.temp_dir, f"temp_{int(time.time())}.tar")
            
            # 使用docker save保存为tar
            cmd = f'docker save {image_name} -o "{temp_tar}"'
            logger.info(f"执行命令: {cmd}")
            
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
            if result.returncode != 0:
                raise Exception(f"Docker save失败: {result.stderr}")
            
            # 使用gzip压缩
            compressed_path = output_path + '.gz'
            logger.info(f"开始压缩镜像文件...")
            
            with open(temp_tar, 'rb') as f_in:
                with gzip.open(compressed_path, 'wb', compresslevel=self.get_compression_level()) as f_out:
                    # 使用大块读取提升性能
                    chunk_size = 16 * 1024 * 1024  # 16MB
                    while True:
                        chunk = f_in.read(chunk_size)
                        if not chunk:
                            break
                        f_out.write(chunk)
            
            # 清理临时文件
            os.unlink(temp_tar)
            
            # 计算压缩比
            original_size = os.path.getsize(temp_tar) if os.path.exists(temp_tar) else 0
            compressed_size = os.path.getsize(compressed_path)
            compression_ratio = (1 - compressed_size / original_size) * 100 if original_size > 0 else 0
            
            logger.info(f"镜像压缩完成，压缩比: {compression_ratio:.1f}%")
            logger.info(f"原始大小: {self._format_size(original_size)}")
            logger.info(f"压缩后大小: {self._format_size(compressed_size)}")
            
            return compressed_path, None
            
        except Exception as e:
            # 清理临时文件
            if os.path.exists(temp_tar):
                os.unlink(temp_tar)
            raise e
    
    def _save_without_compression(self, image_name, output_path):
        """不使用压缩保存镜像"""
        try:
            cmd = f'docker save {image_name} -o "{output_path}"'
            logger.info(f"执行命令: {cmd}")
            
            result = subprocess.run(cmd, shell=True, capture_output=True, text=True)
            if result.returncode != 0:
                raise Exception(f"Docker save失败: {result.stderr}")
            
            return output_path, None
            
        except Exception as e:
            raise e
    
    def transfer_file_optimized(self, local_path, remote_host, remote_path, username, password=None, key_path=None):
        """优化的文件传输"""
        try:
            logger.info(f"开始优化传输文件: {local_path} -> {remote_host}:{remote_path}")
            
            # 根据配置选择传输方式
            if self.get_use_http():
                return self._transfer_via_http(local_path, remote_host, remote_path)
            elif self.get_use_rsync():
                return self._transfer_via_rsync(local_path, remote_host, remote_path, username, key_path)
            else:
                return self._transfer_via_scp_optimized(local_path, remote_host, remote_path, username, password, key_path)
                
        except Exception as e:
            logger.error(f"优化传输失败: {str(e)}")
            return False, str(e)
    
    def _transfer_via_http(self, local_path, remote_host, remote_path):
        """通过HTTP传输文件"""
        try:
            # 端口占用与可用性提示
            port = self.get_http_port()
            if not (1 <= port <= 65535):
                raise Exception(f"HTTP端口无效: {port}")
            
            # 启动HTTP服务器（在后台线程中）
            http_server = self._start_http_server(local_path)
            if http_server is None:
                raise Exception("HTTP服务启动失败，无法进行HTTP传输")
            
            # 构建下载URL
            filename = os.path.basename(local_path)
            download_url = f"http://{remote_host}:{port}/download/{filename}"
            
            logger.info(f"HTTP传输URL: {download_url}")
            
            # 在目标服务器上执行下载命令（调用方负责执行）
            download_cmd = f'wget -O "{remote_path}" "{download_url}" || curl -L -o "{remote_path}" "{download_url}"'
            
            return True, {
                'method': 'http',
                'download_url': download_url,
                'download_cmd': download_cmd,
                'server': http_server
            }
            
        except Exception as e:
            return False, f"HTTP传输准备失败: {str(e)}"

    def _transfer_via_rsync(self, local_path, remote_host, remote_path, username, key_path):
        """通过rsync传输文件"""
        try:
            # 检查rsync是否存在
            which_cmd = 'where' if os.name == 'nt' else 'which'
            check = subprocess.run(f'{which_cmd} rsync', shell=True, capture_output=True, text=True)
            if check.returncode != 0:
                return False, "本机未安装rsync，请安装后再启用或改用SCP/HTTP"
            
            # 构建更适合内网大文件的 rsync 命令：
            # -a 保留属性，-v 详细
            # --whole-file 避免滚动校验（局域网更快）
            # --inplace 直接写目标文件
            # --info=progress2 统一进度输出
            # SSH：关闭压缩，使用较快的加密算法
            ssh_part = f'-e "ssh -T -o Compression=no -c aes128-gcm@openssh.com -c aes128-ctr -i {key_path}"' if key_path else \
                       '-e "ssh -T -o Compression=no -c aes128-gcm@openssh.com -c aes128-ctr"'
            rsync_cmd = f'rsync -av --whole-file --inplace --info=progress2 {ssh_part} "{local_path}" {username}@{remote_host}:{remote_path}'
            
            logger.info(f"执行rsync命令: {rsync_cmd}")
            result = subprocess.run(rsync_cmd, shell=True, capture_output=True, text=True)
            
            if result.returncode == 0:
                logger.info("rsync传输成功")
                return True, None
            else:
                # 提供更清晰的错误提示
                return False, f"rsync传输失败({result.returncode}): {result.stderr or result.stdout}"
                
        except Exception as e:
            return False, f"rsync执行异常: {str(e)}"

    def _transfer_via_scp_optimized(self, local_path, remote_host, remote_path, username, password=None, key_path=None):
        """优化的SCP传输"""
        try:
            # 检查scp命令
            which_cmd = 'where' if os.name == 'nt' else 'which'
            check = subprocess.run(f'{which_cmd} scp', shell=True, capture_output=True, text=True)
            if check.returncode != 0:
                return False, "本机未安装scp，请安装OpenSSH客户端或改用HTTP/rsync"
            
            # 内网大文件：关闭SSH压缩，使用更快算法
            scp_cmd = f'scp -o Compression=no -c aes128-gcm@openssh.com -c aes128-ctr'
            if key_path:
                scp_cmd += f' -i "{key_path}"'
            scp_cmd += f' "{local_path}" {username}@{remote_host}:{remote_path}'
            
            logger.info(f"执行优化SCP命令: {scp_cmd}")
            result = subprocess.run(scp_cmd, shell=True, capture_output=True, text=True)
            
            if result.returncode == 0:
                logger.info("优化SCP传输成功")
                return True, None
            else:
                return False, f"SCP传输失败({result.returncode}): {result.stderr or result.stdout}"
                
        except Exception as e:
            return False, f"SCP执行异常: {str(e)}"
    
    def _start_http_server(self, file_path):
        """启动HTTP服务器用于文件下载"""
        try:
            # 创建简单的HTTP服务器
            server_thread = threading.Thread(
                target=self._run_http_server,
                args=(file_path,),
                daemon=True
            )
            server_thread.start()
            
            # 等待服务器启动
            time.sleep(2)
            
            return server_thread
            
        except Exception as e:
            logger.error(f"启动HTTP服务器失败: {str(e)}")
            return None
    
    def _run_http_server(self, file_path):
        """运行HTTP服务器"""
        try:
            from http.server import HTTPServer, SimpleHTTPRequestHandler
            import socketserver
            
            class CustomHandler(SimpleHTTPRequestHandler):
                def __init__(self, *args, **kwargs):
                    super().__init__(*args, directory=os.path.dirname(file_path), **kwargs)
                
                def do_GET(self):
                    if self.path.startswith('/download/'):
                        filename = self.path[10:]  # 移除 '/download/' 前缀
                        file_path = os.path.join(os.path.dirname(file_path), filename)
                        
                        if os.path.exists(file_path):
                            self.send_response(200)
                            self.send_header('Content-Type', 'application/octet-stream')
                            self.send_header('Content-Disposition', f'attachment; filename="{filename}"')
                            self.end_headers()
                            
                            with open(file_path, 'rb') as f:
                                self.wfile.write(f.read())
                        else:
                            self.send_response(404)
                            self.end_headers()
                    else:
                        super().do_GET()
            
            # 启动服务器
            with socketserver.TCPServer(("", self.get_http_port()), CustomHandler) as httpd:
                logger.info(f"HTTP服务器启动在端口 {self.get_http_port()}")
                httpd.serve_forever()
                
        except Exception as e:
            logger.error(f"HTTP服务器运行失败: {str(e)}")
    
    def parallel_transfer(self, local_path, remote_hosts, remote_paths, username, password=None, key_path=None):
        """并行传输到多个目标"""
        try:
            logger.info(f"开始并行传输到 {len(remote_hosts)} 个目标")
            
            results = []
            
            with ThreadPoolExecutor(max_workers=self.get_parallel_chunks()) as executor:
                # 提交所有传输任务
                future_to_host = {
                    executor.submit(
                        self.transfer_file_optimized,
                        local_path, host, remote_path, username, password, key_path
                    ): host for host, remote_path in zip(remote_hosts, remote_paths)
                }
                
                # 收集结果
                for future in as_completed(future_to_host):
                    host = future_to_host[future]
                    try:
                        success, result = future.result()
                        results.append({
                            'host': host,
                            'success': success,
                            'result': result
                        })
                        
                        if success:
                            logger.info(f"传输到 {host} 成功")
                        else:
                            logger.error(f"传输到 {host} 失败: {result}")
                            
                    except Exception as e:
                        logger.error(f"传输到 {host} 时发生异常: {str(e)}")
                        results.append({
                            'host': host,
                            'success': False,
                            'result': str(e)
                        })
            
            return results
            
        except Exception as e:
            logger.error(f"并行传输失败: {str(e)}")
            return []
    
    def _format_size(self, size_bytes):
        """格式化文件大小"""
        if size_bytes == 0:
            return "0 B"
        size_names = ["B", "KB", "MB", "GB", "TB"]
        i = int(__import__('math').floor(__import__('math').log(size_bytes, 1024)))
        p = __import__('math').pow(1024, i)
        s = round(size_bytes / p, 2)
        return f"{s} {size_names[i]}"
    
    def cleanup_temp_files(self):
        """清理临时文件"""
        try:
            temp_files = [f for f in os.listdir(self.temp_dir) if f.startswith('temp_') and f.endswith('.tar')]
            for temp_file in temp_files:
                file_path = os.path.join(self.temp_dir, temp_file)
                try:
                    os.unlink(file_path)
                    logger.info(f"清理临时文件: {temp_file}")
                except Exception as e:
                    logger.warning(f"清理临时文件失败 {temp_file}: {str(e)}")
        except Exception as e:
            logger.error(f"清理临时文件时发生错误: {str(e)}") 